package com.demo.kafka.processor;

import com.demo.kafka.constants.DelayTopicEnum;
import com.demo.kafka.entity.KafkaMessage;
import com.demo.kafka.producer.KafkaDelayProducer;
import kafka.utils.timer.SystemTimer;
import lombok.extern.slf4j.Slf4j;
import org.apache.kafka.clients.consumer.Consumer;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.ConsumerRecords;

import java.time.Duration;
import java.util.*;
import java.util.stream.Collectors;

/**
 * 延迟消息处理
 * @author K
 */
@Slf4j
public class KafkaDelayMsgProcessor {

    private KafkaDelayProducer delayProducer;
    private Map<String, Consumer<String, KafkaMessage>> consumerMap;

    public KafkaDelayMsgProcessor(KafkaDelayProducer delayProducer, Map<String, Consumer<String, KafkaMessage>> consumerMap) {
        this.delayProducer = delayProducer;
        this.consumerMap = consumerMap;
    }

    /**
     * 获取分区延迟时间，并且计算延迟消息是否到达发送时间
     * @param topic 主题
     * @return 延迟时间
     */
    public void processDelayMsgAndGetPartitionDelayTime(String topic) {
        log.debug("consumer poll message");
        Consumer<String, KafkaMessage> consumer = consumerMap.get(topic);
        ConsumerRecords<String, KafkaMessage> consumerRecords = consumer.poll(Duration.ofMillis(200));
        log.info("主题：{}， consumerRecords count ： {}", topic, consumerRecords.count());
        if (consumerRecords.isEmpty()) {
            return;
        }

        Iterable<ConsumerRecord<String, KafkaMessage>> records = consumerRecords.records(topic);
        for (ConsumerRecord<String, KafkaMessage> record : records) {
            KafkaMessage kafkaMessage = record.value();
            long currentTimeMillis = System.currentTimeMillis();
            long msgTimestamp = record.timestamp();
            long msgDelayTime = kafkaMessage.getDelayTime();
            log.info("主题：{}，分区：{}，当前时间：{}， 发送消息时间 : {}， 延时时间：{}",
                    record.topic(), record.partition(), new Date(currentTimeMillis), new Date(msgTimestamp),
                    msgDelayTime);

            // 计算消息剩于的发送时间，公式：发送消息时间+延迟时间-当前时间
            long diff = msgTimestamp + msgDelayTime - currentTimeMillis;
            // 差值小于500毫秒不再走一次时间轮，立刻发送消息到指定的topic
            if (diff <= 500) {
                delayProducer.sendDelay(kafkaMessage.getTopic(), kafkaMessage.getPayload(), 0);
                log.info("diff:{}, 消息直接投递，topic:{}", diff, kafkaMessage.getTopic());
            } else if (diff <= 1000) {
                delayProducer.sendDelay(kafkaMessage.getTopic(), kafkaMessage.getPayload(), 1000);
                log.info("diff:{}, 发送延迟1秒消息，topic:{}", diff, kafkaMessage.getTopic());
            } else {
                delayProducer.sendDelay(kafkaMessage.getTopic(), kafkaMessage.getPayload(), diff);
                log.info("发送延迟{}秒消息，topic:{}", diff, kafkaMessage.getTopic());
            }
        }
        consumer.commitSync();
    }

    /**
     * 分区延迟执行事件添加到时间轮里
     * @param queue 到期发送事件队列
     */
    public void addTimeTask(Queue<String> queue) {
        List<DelayTopicEnum> list = Arrays.asList(DelayTopicEnum.values());
        addTimeTask(list, queue);
    }

    public void addTimeTask(List<DelayTopicEnum> list, Queue<String> queue) {
        // 初始化任务
        List<KafkaTimerTaskProcessor> timerTasks = initTimerTasks(list, queue);
        SystemTimer systemTimer = KafkaTimerTaskProcessor.getSystemTimer();
        for (KafkaTimerTaskProcessor timerTask : timerTasks) {
            systemTimer.add(timerTask);
        }
    }

    private List<KafkaTimerTaskProcessor> initTimerTasks(List<DelayTopicEnum> list, Queue<String> queue) {

        return list.stream().map(delayTopicEnum ->
            new KafkaTimerTaskProcessor(delayTopicEnum.getTopic(), delayTopicEnum.getTime(), queue)
        ).collect(Collectors.toList());
    }
}
